import torch  # 导入PyTorch库

# 定义模型参数
input_size = 4  # 输入词汇表大小（字符总数）
num_class = 4  # 输出类别数（与输入词汇表大小相同）
hidden_size = 8  # RNN隐藏层维度
embedding_size = 10  # 嵌入层输出维度（将字符索引映射到10维向量）
batch_size = 1  # 批次大小
num_layers = 2  # RNN层数（这里使用2层）
seq_len = 5  # 序列长度

# 字符到索引的映射表
idx2char_1 = ['e', 'h', 'l', 'o']  # 完整字符集：0->e, 1->h, 2->l, 3->o
idx2char_2 = ['h', 'l', 'o']  # 备用字符集（未在代码中使用）

# 输入输出数据
x_data = [[1, 0, 2, 2, 3]]  # 输入序列索引：对应"h e l l o"，形状为(batch_size, seq_len)
y_data = [3, 1, 2, 3, 2]  # 目标序列索引：对应"o h l o l"

# 转换为PyTorch张量
inputs = torch.LongTensor(x_data)  # 输入张量，形状为(batch_size, seq_len)
labels = torch.LongTensor(y_data)  # 标签张量，形状为(batch_size*seq_len,)


# 定义模型类
class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()

        # 嵌入层：将整数索引映射为稠密向量
        # input_size：词汇表大小
        # embedding_size：输出向量维度
        self.emb = torch.nn.Embedding(input_size, embedding_size)

        # RNN层：使用两层RNN
        # input_size：输入特征维度（这里是嵌入层输出维度）
        # hidden_size：隐藏层维度
        # num_layers：RNN层数
        # batch_first=True：输入输出张量的第一个维度为batch_size
        self.rnn = torch.nn.RNN(
            input_size=embedding_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True
        )

        # 全连接层：将RNN输出映射到类别空间
        self.fc = torch.nn.Linear(hidden_size, num_class)

    def forward(self, x):
        # 初始化隐藏状态：形状为(num_layers, batch_size, hidden_size)
        hidden = torch.zeros(num_layers, x.size(0), hidden_size)

        # 嵌入层处理：将索引转换为稠密向量
        # 输入x形状：(batch_size, seq_len)
        # 输出x形状：(batch_size, seq_len, embedding_size)
        x = self.emb(x)

        # RNN前向计算
        # 输出x形状：(batch_size, seq_len, hidden_size)
        x, _ = self.rnn(x, hidden)

        # 全连接层处理：映射到类别空间
        # 输出x形状：(batch_size, seq_len, num_class)
        x = self.fc(x)

        # 重塑输出形状：(batch_size*seq_len, num_class)，适应交叉熵损失计算
        return x.view(-1, num_class)


# 创建模型实例
net = Model()

# 定义损失函数和优化器
criterion = torch.nn.CrossEntropyLoss()  # 交叉熵损失，适用于分类问题
optimizer = torch.optim.Adam(net.parameters(), lr=0.05)  # Adam优化器

# 训练循环
for epoch in range(15):  # 训练15轮
    optimizer.zero_grad()  # 清零梯度

    # 前向传播：获取模型输出
    outputs = net(inputs)  # outputs形状：(batch_size*seq_len, num_class)

    # 计算损失
    loss = criterion(outputs, labels)

    # 反向传播与参数更新
    loss.backward()  # 计算梯度
    optimizer.step()  # 更新参数

    # 计算预测结果（取概率最大的索引）
    _, idx = outputs.max(dim=1)  # idx是预测的字符索引
    idx = idx.data.numpy()  # 转换为numpy数组

    # 打印预测结果和损失
    print('Predicted string: ', ''.join([idx2char_1[x] for x in idx]), end='')
    print(", Epoch [%d/15] loss = %.3f" % (epoch + 1, loss.item()))